from com.example.tree.tree_node import TreeNode
from typing import Optional, List


def beforeOrderTraversal(root: Optional[TreeNode], res: List[int]) -> None:
    """
    前序遍历(迭代实现)
    :param res:
    :param root:
    :return:
    """
    # 写法一(与中序、后序类似 模板解法)
    stack = []
    tmp = root
    while tmp or stack:
        while tmp:
            res.append(tmp.val)  # 访问根节点
            stack.append(tmp)
            tmp = tmp.left  # 访问左节点
        tmp = stack.pop()
        tmp = tmp.right  # 访问右节点

    # 写法二:
    # if not root:
    #     return
    # stack = [root]
    # while stack:
    #     tmp = stack.pop()
    #     res.append(tmp.val)
    #     if tmp.right:
    #         stack.append(tmp.right)
    #     if tmp.left:
    #         stack.append(tmp.left)


def inOrderTraversal(root: Optional[TreeNode], res: List[int]) -> None:
    """
    中序遍历(迭代实现)
    :param res:
    :param root:
    :return:
    """
    # 写法一:模板解法
    stack = []
    tmp = root
    while tmp or stack:
        while tmp:  # 一直向左直到最底
            stack.append(tmp)
            tmp = tmp.left
        tmp = stack.pop()  # 从栈中弹出的节点，其左子树肯定已经是访问完了的，因此可以直接访问该节点，然后接下来访问右子树
        res.append(tmp.val)  # 访问根节点
        tmp = tmp.right  # 访问右节点

    # 写法二:
    # if not root:
    #     return res
    # tmp, stack = root, []
    # while tmp:
    #     if tmp.left:  # 存在左子树，说明该树的根节点还不能访问就将节点入栈并深入左子树
    #         stack.append(tmp)
    #         tmp = tmp.left
    #     else:  # 否则就可以寻找深入右子树的节点
    #         while stack and not tmp.right:
    #             res.append(tmp.val)
    #             tmp = stack.pop()
    #         res.append(tmp.val)  # 如果可以深入右子树，访问该节点后，深入右子树
    #         tmp = tmp.right


def inOrderTraversal2(root: Optional[TreeNode], res: List[int]) -> None:
    """
    Morris 中序遍历
    将非递归的中序遍历空间复杂度降为 O(1)，但是每个节点会被访问两次
    改变了整个树的结构，强行把一棵二叉树改成一段链表结构(第二次遍历时恢复)
    Morris 遍历其本质是线索二叉树，就是利用二叉树中 n+1 个指向None的指针。
    在遍历的过程中，通过利用叶子节点空的right指针，指向中序遍历的后继节点，从而避免了对 stack 的依赖。

    具体实现：
    将遍历到的当前节点记为tmp
    如果 tmp 无左孩子，先将 tmp 的值加入答案数组，再访问 tmp 的右孩子，即 tmp=tmp.right。
    如果 tmp 有左孩子，则找到 tmp 左子树上最右的节点（即左子树中序遍历的最后一个节点，tmp 在中序遍历中的前驱节点），记为 pre(原树的其右子节点一定为空)。
    根据 pre 的右孩子是否为空，进行如下操作。
        如果 pre 的右孩子为空，则将其右孩子指向 tmp，然后访问 tmp 的左孩子，即 tmp = tmp.left。
        如果 pre 的右孩子不为空，则此时其右孩子指向 tmp，说明已经遍历完 tmp 的左子树，
        将 pre 的右孩子置空，将 tmp 的值加入答案数组，然后访问 tmp 的右孩子，即 tmp=tmp.right。

    :param root:
    :return:
    """
    tmp, pre = root, None
    while tmp:
        if tmp.left:  # 当前节点有左子树，将当前节点连带右子树全部挂到左节点的最右子树下面
            # pre 节点就是当前 tmp 节点向左走一步，然后一直向右走至无法走为止
            pre = tmp.left
            while pre.right and pre.right is not tmp:
                pre = pre.right
            if pre.right:  # 是第二次访问了，那么应当消除线索还原原来的二叉树，将节点的值加入到res中，然后访问其右子树
                res.append(tmp.val)
                pre.right = None
                tmp = tmp.right
            else:  # 是第一次访问，建立线索
                pre.right = tmp
                tmp = tmp.left
        else:  # 左子树为空，则打印这个节点，并向右边遍历
            res.append(tmp.val)
            tmp = tmp.right
    return res


def postorderTraversal(root: Optional[TreeNode], res: List[int]) -> None:
    """
    后续遍历(迭代实现) 左 -> 右 -> 根
    :param res:
    :param root:
    :return:
    """
    stack = []
    tmp = root
    # 每次出栈root不断上移，下一次循环时还会判断是否有右子树；如果不加prev就会陷入死循环~，加个prev做个判断，只有没有遍历过的右子树才能入栈
    pre = None  # 记录访问历史，在回溯到父节点时，可以由此来判断，上一个访问的节点是否为右子树
    while tmp or stack:
        while tmp:
            stack.append(tmp)
            tmp = tmp.left
        tmp = stack.pop()  # 从栈中弹出的节点，只能确定其左子树肯定访问完了，但是无法确定右子树是否访问过
        if not tmp.right or tmp.right == pre:  # 如果没有右子树，或者右子树访问完了，也就是上一个访问的节点是右子节点时可以访问当前节点
            res.append(tmp.val)
            pre = tmp  # 更新历史访问记录，这样回溯的时候父节点可以由此判断右子树是否访问完成
            tmp = None
        else:  # 如果右子树没有被访问，那么将当前节点压栈，访问右子树
            stack.append(tmp)
            tmp = tmp.right


if __name__ == "__main__":
    #     3
    #    / \
    #   9  20
    #     /  \
    #    15   7
    root = TreeNode(3)
    root.left, root.right = TreeNode(9), TreeNode(20)
    root.right.left, root.right.right = TreeNode(15), TreeNode(7)
    print("==============前序遍历(迭代实现)===============")
    lst1 = []
    beforeOrderTraversal(root, lst1)
    print(lst1)
    print("==============中序遍历(迭代实现)===============")
    lst2 = []
    inOrderTraversal(root, lst2)
    print(lst2)
    print("==============后序遍历(迭代实现)===============")
    lst3 = []
    postorderTraversal(root, lst3)
    print(lst3)
